import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os
import warnings
warnings.filterwarnings('ignore')

class Ic_record():
    def __init__(self, args, codes) -> None:
        self.args = args
        self.codes = codes   
        self.test_recall_list = [[0 for i in range(len(codes))]]
        self.test_precision_list = [[0 for i in range(len(codes))]]
        self.test_f1_list = [[0 for i in range(len(codes))]]
        self.train_recall_list = []
        self.train_precision_list = []
        self.train_f1_list = []
        self.eval_precision_list= []
        self.eval_recall_list= []
        self.eval_f1_list= []
        
    def ic_record(self, pred, true, root, times, name='test'):  
        pred_df, true_df, recall_list, precision_list, f1_list = self.array_to_df(pred, true, times)
        if name == 'test':
            self.test_precision_list.append(precision_list)
            self.test_recall_list.append(recall_list)
            self.test_f1_list.append(f1_list)
        if name == 'eval':
            self.eval_precision_list.append(precision_list)
            self.eval_recall_list.append(recall_list)
            self.eval_f1_list.append(f1_list)
        if name == 'train':
            self.train_precision_list.append(precision_list)
            self.train_recall_list.append(recall_list)
            self.train_f1_list.append(f1_list)          
        print( f'当前{name}f1_mean：{np.nanmean(f1_list)}, f1_std：{np.nanstd(f1_list)},precsion：{np.nanmean(precision_list)}')
        if self.args.predict_save:
            pred_df.to_pickle(os.path.join(root, f'{name}_pred.pkl.gzip'))
            true_df.to_pickle(os.path.join(root, f'{name}_true.pkl.gzip'))
        plot_line_ml(array1=pred_df.iloc[-1000:,0].values, array2=true_df.iloc[-1000:,0].values,   name=name, root=root) #times=times,
         
    def save_ic(self, root, month_index):
        test_recall_df = ic_re_formate(self.test_recall_list, month_index, self.codes)
        test_recall_df.to_csv(f'{root}/test_recall.csv')
        eval_recall_df = ic_re_formate(self.eval_recall_list, month_index, self.codes)
        eval_recall_df.to_csv(f'{root}/eval_recall.csv')
        train_recall_df = ic_re_formate(self.train_recall_list, month_index, self.codes)
        train_recall_df.to_csv(f'{root}/train_recall.csv')

        test_precision_df = ic_re_formate(self.test_precision_list, month_index, self.codes)
        test_precision_df.to_csv(f'{root}/test_precision.csv')
        eval_precision_df = ic_re_formate(self.eval_precision_list, month_index, self.codes)
        eval_precision_df.to_csv(f'{root}/eval_precision.csv')
        train_precision_df = ic_re_formate(self.train_precision_list, month_index, self.codes)
        train_precision_df.to_csv(f'{root}/train_precision.csv')

        test_f1_df = ic_re_formate(self.test_f1_list, month_index, self.codes)
        test_f1_df.to_csv(f'{root}/test_f1.csv')
        eval_f1_df = ic_re_formate(self.eval_f1_list, month_index, self.codes)
        eval_f1_df.to_csv(f'{root}/eval_f1.csv')
        train_f1_df = ic_re_formate(self.train_f1_list, month_index, self.codes)
        train_f1_df.to_csv(f'{root}/train_f1.csv')
        # test_rmse_df = ic_re_formate(self.test_precision_list, month_index, self.codes)
        # train_rmse_df = ic_re_formate(self.train_precision_list, month_index, self.codes)

    def array_to_df(self, pred, true, eval_times):
        pred = shape_squeeze(pred)
        true = shape_squeeze(true)
        pred_df = pd.DataFrame(pred, index=eval_times[self.args.seq_len:], columns=self.codes)#index_time[:len(pred)]
        true_df = pd.DataFrame(true, index=eval_times[self.args.seq_len:], columns=self.codes)
        recall_list = []
        precision_list = []
        f1_list = []
        TPR, precision, f1 = calc_metric(pred, true)
        print( f'f1：{f1}, precision：{precision},recall：{TPR}')
        for i in range(pred.shape[1]):
            TPR, precision, f1 = calc_metric(pred[:, i], true[:, i])
            recall_list.append(TPR)
            precision_list.append(precision)
            f1_list.append(f1)   
        return pred_df, true_df, recall_list, precision_list, f1_list


def ic_re_formate(recall_list, month_index, codes):
    ic = np.array(recall_list)
    ic_df = pd.DataFrame(ic, index=month_index, columns=codes)
    ic_values = ic_df.values
    if len(ic_values.shape)>1:
        mean_id = np.nanmean(ic_values, axis = 0)
        std_id = np.nanstd(ic_values, axis = 0)
        lst = [mean_id, std_id]
        id_df = pd.DataFrame(lst, index=['mean', 'std'], columns=ic_df.columns)
        ic_df = pd.concat([ic_df, id_df], axis=0)
        ic_values_id = ic_df.values 
        mean_id = np.nanmean(ic_values_id, axis = 1)
        std_id = np.nanstd(ic_values_id, axis = 1)
        ic_df['mean'] = mean_id
        ic_df['std'] = std_id
    return ic_df

def calc_metric(pred, true):
    pred = (pred>=0.6).astype(int)
    # pair = ((pred==true).astype(int))
    # TP = pair[pred==1]
    TP = ((pred==1)&(true==1)).astype(int)
    FN = ((pred==0)&(true==1)).astype(int)

    FP = ((pred==1)&(true==0)).astype(int)
    TN = ((pred==0)&(true==0)).astype(int)

    TPR = np.nansum(TP)/(np.nansum(TP)+np.nansum(FN)+0.00000001)
    FPR = np.nansum(FP)/(np.nansum(FP)+np.nansum(TN)+0.00000001)
    precision = np.nansum(TP)/(np.nansum(TP)+np.nansum(FP)+0.00000001)
    f1 = 2*precision*TPR/(precision+TPR)
    if len(TP.shape)>1:
        print('p ration:', np.nansum(TP)/(TP.shape[0]*TP.shape[1]), np.nansum(TP))
    return TPR, precision, f1

def shape_squeeze(data):
    shape = data.shape  
    if len(shape)>2:
        data = np.squeeze(data, axis=-1)
    return data

def plot_line_ml(array1, array2, name, root):
    plt.figure(figsize=(30, 10))
    plt.plot(array1,  alpha=0.5, label='pred')
    plt.plot(array2, alpha=0.5, label='true')
    plt.plot(np.ones(len(array1))*0.1, label = '0.1' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.2, label = '0.2' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.5, label = '0.5' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.7, label = '0.7' , alpha=0.5, linestyle=':')
    plt.legend()
    save_root = f'{root}/{name}_result.png'
    plt.savefig(save_root)
    plt.close()